package se.haaga.jsupervision.monitor.linux;

import java.awt.Container;
import java.awt.GridBagConstraints;
import java.awt.GridBagLayout;
import java.awt.Insets;
import java.io.IOException;
import java.io.InputStream;
import java.math.BigInteger;
import java.util.HashMap;
import java.util.Map;

import javax.swing.JLabel;
import javax.swing.JOptionPane;
import javax.swing.JPanel;
import javax.swing.JPasswordField;
import javax.swing.JTextField;

import com.jcraft.jsch.Channel;
import com.jcraft.jsch.ChannelExec;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import com.jcraft.jsch.UIKeyboardInteractive;
import com.jcraft.jsch.UserInfo;

public class LinuxMonitor {
	
	private final String host;
	
	private final String username;
	
	private final int port;
	
	private boolean connectionActive = false;
	
	private Session session;
	
	public enum LinuxVariables {
		SYS_MEM_TOTAL, SYS_MEM_AVAILABLE, SWAP_TOTAL, SWAP_AVAILABLE,
		DISTRIBUTION_ID, DISTRIBUTION_RELEASE, DISTRIBUTION_CODENAME, DISTRIBUTION_DESCRIPTION,
		KERNEL_VERSION, UPTIME, TOTAL_CPUS_FOUND, CPU_LOAD, DRIVE_SPACE
	}
	
	public LinuxMonitor(String host, String username, int port) {
		this.host = host;
		this.username = username;
		this.port = port;
	}
	
	public void closeConnection() {
		if (session != null && session.isConnected()) {
			session.disconnect();
		}
	}
	
	public void openConnection() {
		try {
			JSch jsch = new JSch();
			session = jsch.getSession(username, host, port);
			UserInfo ui = new MyUserInfo();
			session.setUserInfo(ui);
			session.connect();
		} catch (JSchException e) {
			System.err.println(e);
		}
	}
	
	public Map<LinuxVariables, String> getCPULoad() {
		String cpu = executeCommand("cat /proc/stat | grep '^cpu'");
		return fetchCPUData(cpu);
	}
	
	private Map<LinuxVariables, String> fetchCPUData(String cpuInfo) {
		Map<LinuxVariables, String> returnMap = new HashMap<LinuxVariables, String>();
		if (cpuInfo != null) {
			int totalCpus = 0;
			while (true) {
				if (cpuInfo.contains("cpu" + totalCpus)) {
					totalCpus++;
				} else {
					break;
				}
			}
			returnMap.put(LinuxVariables.TOTAL_CPUS_FOUND, String.valueOf(totalCpus));
			int lastCpu = totalCpus - 1;
			Map<String, BigInteger> cpuData = new HashMap<String, BigInteger>();
			for (int i = 0; i < totalCpus; i++) {
				BigInteger u1;
				BigInteger n1;
				BigInteger s1;
				BigInteger i1;
				String tmp;
				if (i == lastCpu) {
					tmp = cpuInfo.substring(cpuInfo.indexOf("cpu" + i), cpuInfo.length());
				} else {
					tmp = cpuInfo.substring(cpuInfo.indexOf("cpu" + i), cpuInfo.indexOf("cpu" + (i + 1)));
				}
				tmp = tmp.replaceAll("cpu" + i, "");
				tmp = tmp.trim();
				String[] splitted = tmp.split(" ");
				u1 = new BigInteger(splitted[0]);
				n1 = new BigInteger(splitted[1]);
				s1 = new BigInteger(splitted[2]);
				i1 = new BigInteger(splitted[3]);
				cpuData.put("cpu" + i + "_u1", u1);
				cpuData.put("cpu" + i + "_n1", n1);
				cpuData.put("cpu" + i + "_s1", s1);
				cpuData.put("cpu" + i + "_i1", i1);
			}
			try {
				Thread.sleep(50);
			} catch (InterruptedException ex) {
				ex.printStackTrace();
			}
			// Fetch again
			cpuInfo = executeCommand("cat /proc/stat | grep '^cpu'");
			for (int i = 0; i < totalCpus; i++) {
				BigInteger u2;
				BigInteger n2;
				BigInteger s2;
				BigInteger i2;
				String tmp;
				if (i == lastCpu) {
					tmp = cpuInfo.substring(cpuInfo.indexOf("cpu" + i), cpuInfo.length());
				} else {
					tmp = cpuInfo.substring(cpuInfo.indexOf("cpu" + i), cpuInfo.indexOf("cpu" + (i + 1)));
				}
				tmp = tmp.replaceAll("cpu" + i, "");
				tmp = tmp.trim();
				String[] splitted = tmp.split(" ");
				u2 = new BigInteger(splitted[0]);
				n2 = new BigInteger(splitted[1]);
				s2 = new BigInteger(splitted[2]);
				i2 = new BigInteger(splitted[3]);
				cpuData.put("cpu" + i + "_u2", u2);
				cpuData.put("cpu" + i + "_n2", n2);
				cpuData.put("cpu" + i + "_s2", s2);
				cpuData.put("cpu" + i + "_i2", i2);
			}
			// Now, calculate CPU usage for each cpu
			Map<String, String> tmpMap = new HashMap<String, String>();
			for (int i = 0; i < totalCpus; i++) {
				BigInteger u1 = cpuData.get("cpu" + i + "_u1");
				BigInteger n1 = cpuData.get("cpu" + i + "_n1");
				BigInteger s1 = cpuData.get("cpu" + i + "_s1");
				BigInteger i1 = cpuData.get("cpu" + i + "_i1");
				BigInteger u2 = cpuData.get("cpu" + i + "_u2");
				BigInteger n2 = cpuData.get("cpu" + i + "_n2");
				BigInteger s2 = cpuData.get("cpu" + i + "_s2");
				BigInteger i2 = cpuData.get("cpu" + i + "_i2");
				BigInteger u12 = u2.subtract(u1);
				BigInteger n12 = n2.subtract(n1);
				BigInteger s12 = s2.subtract(s1);
				BigInteger i12 = i2.subtract(i1);
				BigInteger usage = u12.add(n12);
				usage = usage.add(s12);
				BigInteger total = usage.add(i12);
				BigInteger tmp = usage.multiply(new BigInteger("100"));
				String result = tmp.divide(total).toString();
				tmpMap.put("cpu" + i, result);
			}
			String s = "";
			for (int i = 0; i < totalCpus; i++) {
				if (lastCpu == i) {
					s += "cpu" + i + ";" + tmpMap.get("cpu" + i);
				} else {
					s += "cpu" + i + ";" + tmpMap.get("cpu" + i) + "@";
				}
			}
			returnMap.put(LinuxVariables.CPU_LOAD, s);
		}
		return returnMap;
	}

	public Map<LinuxVariables, Integer> getMemoryStatus() {
		String mem = executeCommand("free -b");
		return parseMemData(mem);
	}
	
	public Map<LinuxVariables, String> getSystemInformation() {
		String dist = executeCommand("cat /etc/*release");
		Map<LinuxVariables, String> distData = parseDistData(dist);
		String kernel = executeCommand("uname -r");
		distData.put(LinuxVariables.KERNEL_VERSION, kernel);
		String uptime = executeCommand("uptime");
		distData.put(LinuxVariables.UPTIME, parseUptimeString(uptime));
		return distData;
	}
	
	private String parseUptimeString(String uptime) {
		String returnString = null;
		if (uptime != null) {
			String s = uptime.substring(uptime.indexOf("up"), uptime.indexOf(","));
			s = s.replaceAll("up", "");
			returnString = s.trim();
		}
		return returnString;
	}

	private Map<LinuxVariables, String> parseDistData(String dist) {
		Map<LinuxVariables, String> returnMap = new HashMap<LinuxVariables, String>();
		if (dist != null) {
			String id = dist.substring(dist.indexOf("DISTRIB_ID="), dist.indexOf("DISTRIB_RELEASE="));
			String release = dist.substring(dist.indexOf("DISTRIB_RELEASE="), dist.indexOf("DISTRIB_CODENAME="));
			String codename = dist.substring(dist.indexOf("DISTRIB_CODENAME="), dist.indexOf("DISTRIB_DESCRIPTION="));
			String desc = dist.substring(dist.indexOf("DISTRIB_DESCRIPTION="), dist.length());

			id = id.replaceAll("DISTRIB_ID=", "").trim();
			release = release.replaceAll("DISTRIB_RELEASE=", "").trim();
			codename = codename.replaceAll("DISTRIB_CODENAME=", "").trim();
			desc = desc.replaceAll("DISTRIB_DESCRIPTION=", "").replaceAll("\"", "").trim();
			returnMap.put(LinuxVariables.DISTRIBUTION_ID, id);
			returnMap.put(LinuxVariables.DISTRIBUTION_RELEASE, release);
			returnMap.put(LinuxVariables.DISTRIBUTION_CODENAME, codename);
			returnMap.put(LinuxVariables.DISTRIBUTION_DESCRIPTION, desc);
		}
		return returnMap;
	}

	private Map<LinuxVariables, Integer> parseMemData(String memData) {
		Map<LinuxVariables, Integer> returnMap = new HashMap<LinuxVariables, Integer>();
		if (memData != null) {
			memData = memData.substring(memData.indexOf("Mem: "), memData.length());
			String memTemp = memData.substring(memData.indexOf("Mem: "), memData.indexOf("-/+"));
			memTemp = memTemp.replaceAll("Mem: ", "");
			memTemp = memTemp.trim();
			String[] memArray = new String[3];
			memArray[0] = "sysmem";
			String tmp = "";
			for (int i = 0; i < memTemp.length(); i++) {
				char current = memTemp.charAt(i);
				String currentString = String.valueOf(current);
				if (!tmp.equals("") && currentString.equals(" ")) {
					if (memArray[1] == null) {
						memArray[1] = tmp;
					} else if (memArray[2] == null) {
						memArray[2] = tmp;
						break;
					}
				}
				if (!currentString.equals(" ")) {
					tmp += currentString;
				} else {
					tmp = "";
				}
			}
			returnMap.put(LinuxVariables.SYS_MEM_TOTAL, new BigInteger(memArray[1]).divide(new BigInteger("1000000")).intValue());
			returnMap.put(LinuxVariables.SYS_MEM_AVAILABLE, new BigInteger(memArray[2]).divide(new BigInteger("1000000")).intValue());
			memData = memData.substring(memData.indexOf("Mem: "), memData.length());
			String swapTemp = memData.substring(memData.indexOf("Swap: "), memData.length());
			swapTemp = swapTemp.replaceAll("Swap: ", "");
			swapTemp = swapTemp.trim();
			String[] swapArray = new String[3];
			swapArray[0] = "swap";
			String tmpSwapHolder = "";
			for (int i = 0; i < swapTemp.length(); i++) {
				char current = swapTemp.charAt(i);
				String currentString = String.valueOf(current);
				if (!tmpSwapHolder.equals("") && currentString.equals(" ")) {
					if (swapArray[1] == null) {
						swapArray[1] = tmpSwapHolder;
					} else if (swapArray[2] == null) {
						swapArray[2] = tmpSwapHolder;
						break;
					}
				}
				if (!currentString.equals(" ")) {
					tmpSwapHolder += currentString;
				} else {
					tmpSwapHolder = "";
				}
			}
			returnMap.put(LinuxVariables.SWAP_TOTAL, new BigInteger(swapArray[1]).divide(new BigInteger("1000000")).intValue());
			returnMap.put(LinuxVariables.SWAP_AVAILABLE, new BigInteger(swapArray[2]).divide(new BigInteger("1000000")).intValue());
		}
		return returnMap;
	}
	
	public Map<LinuxVariables, String> getDriveInfo() {
		String driveInfo = executeCommand("/bin/df -Pm |awk '{print $3  \" \" $4 \" \" $5  \" \" $6}'");
		return parseDriveInfo(driveInfo);
	}
	
	private Map<LinuxVariables, String> parseDriveInfo(String driveInfo) {
		Map<LinuxVariables, String> returnMap = new HashMap<LinuxVariables, String>();
		if (driveInfo != null) {
			driveInfo = driveInfo.substring(driveInfo.indexOf("Mounted") + "Mounted".length());
			driveInfo = driveInfo.trim();
			driveInfo = driveInfo.replaceAll("\n", " ");
			String[] split = driveInfo.split(" ");
			String returnString = "";
			for (int i = 0; i < split.length; i++) {
				if (split[i] != null) {
					if (i + 1 < split.length && i + 2 < split.length && i + 3 < split.length) {
						returnString += split[i] + ";";
						returnString += split[++i] + ";";
						returnString += split[++i] + ";";
						returnString += split[++i] + "@";
					}
				}
			}
			returnString = returnString.substring(0, returnString.length() - 1);
			returnMap.put(LinuxVariables.DRIVE_SPACE, returnString);
		}
		return returnMap;
	}

	public boolean isConnectionActive() {
		return connectionActive;
	}
	
	private String executeCommand(String command) {
		String returnString = "";
		try {
			Channel channel = session.openChannel("exec");
			((ChannelExec) channel).setCommand(command);
	
			channel.setInputStream(null);
	
			((ChannelExec) channel).setErrStream(System.err);
	
			InputStream in = channel.getInputStream();
	
			channel.connect();
	
			byte[] tmp = new byte[1024];
			while (true) {
				while (in.available() > 0) {
					int i = in.read(tmp, 0, 1024);
					if (i < 0)
						break;
					returnString += (new String(tmp, 0, i));
				}
				if (channel.isClosed()) {
					break;
				}
				try {
					Thread.sleep(1000);
				} catch (Exception e) {
					System.err.println(e);
				}
			}
			channel.disconnect();
		} catch (JSchException e) {
			System.out.println(e);
		} catch (IOException e) {
			System.err.println(e);
		}
		return returnString;
	}
	
	public static class MyUserInfo implements UserInfo, UIKeyboardInteractive {
		
		public String getPassword() {
			return passwd;
		}

		public boolean promptYesNo(String str) {
			Object[] options = { "yes", "no" };
			int foo = JOptionPane.showOptionDialog(null, str, "Warning",
					JOptionPane.DEFAULT_OPTION, JOptionPane.WARNING_MESSAGE,
					null, options, options[0]);
			return foo == 0;
		}

		String passwd;
		
		JTextField passwordField = (JTextField) new JPasswordField(20);

		public String getPassphrase() {
			return null;
		}

		public boolean promptPassphrase(String message) {
			return true;
		}

		public boolean promptPassword(String message) {
			Object[] ob = { passwordField };
			int result = JOptionPane.showConfirmDialog(null, ob, message,
					JOptionPane.OK_CANCEL_OPTION);
			if (result == JOptionPane.OK_OPTION) {
				passwd = passwordField.getText();
				return true;
			} else {
				return false;
			}
		}

		public void showMessage(String message) {
			JOptionPane.showMessageDialog(null, message);
		}

		final GridBagConstraints gbc = new GridBagConstraints(0, 0, 1, 1, 1, 1,
				GridBagConstraints.NORTHWEST, GridBagConstraints.NONE,
				new Insets(0, 0, 0, 0), 0, 0);
		private Container panel;

		public String[] promptKeyboardInteractive(String destination,
				String name, String instruction, String[] prompt, boolean[] echo) {
			panel = new JPanel();
			panel.setLayout(new GridBagLayout());

			gbc.weightx = 1.0;
			gbc.gridwidth = GridBagConstraints.REMAINDER;
			gbc.gridx = 0;
			panel.add(new JLabel(instruction), gbc);
			gbc.gridy++;

			gbc.gridwidth = GridBagConstraints.RELATIVE;

			JTextField[] texts = new JTextField[prompt.length];
			for (int i = 0; i < prompt.length; i++) {
				gbc.fill = GridBagConstraints.NONE;
				gbc.gridx = 0;
				gbc.weightx = 1;
				panel.add(new JLabel(prompt[i]), gbc);

				gbc.gridx = 1;
				gbc.fill = GridBagConstraints.HORIZONTAL;
				gbc.weighty = 1;
				if (echo[i]) {
					texts[i] = new JTextField(20);
				} else {
					texts[i] = new JPasswordField(20);
				}
				panel.add(texts[i], gbc);
				gbc.gridy++;
			}

			if (JOptionPane.showConfirmDialog(null, panel, destination + ": "
					+ name, JOptionPane.OK_CANCEL_OPTION,
					JOptionPane.QUESTION_MESSAGE) == JOptionPane.OK_OPTION) {
				String[] response = new String[prompt.length];
				for (int i = 0; i < prompt.length; i++) {
					response[i] = texts[i].getText();
				}
				return response;
			} else {
				return null; // cancel
			}
		}
	}
	
}
